package com.xlh.dokka.svc.core.scheduler.strategy;

import com.xlh.dokka.svc.core.model.ContainerConfig;
import com.xlh.dokka.svc.core.model.GpuNode;
import com.xlh.dokka.svc.core.model.Node;
import org.apache.commons.collections.CollectionUtils;

import java.util.ArrayList;
import java.util.Comparator;
import java.util.List;

/**
 * @author cheer
 */
public class ResourcesSpreadStrategy implements Strategy {

    @Override
    public List<Node> rankAndSort(ContainerConfig config, List<Node> nodes) {
        List<Node> result = new ArrayList<>();
        if (CollectionUtils.isNotEmpty(nodes)) {
            List<WeightedNode> weightedNodes = new ArrayList<>();
            for (Node node : nodes) {
                WeightedNode weightedNode = new WeightedNode();
                weightedNode.setNode(node);
                weightedNode.setScore(computeScore(config, node));
                weightedNodes.add(weightedNode);
            }
            weightedNodes.sort(Comparator.comparingDouble(WeightedNode::getScore));
            for (WeightedNode weightedNode : weightedNodes) {
                result.add(weightedNode.getNode());
            }
        }
        return result;
    }

    @Override
    public List<GpuNode> rankAndSortGpu(ContainerConfig config, List<GpuNode> gpuNodes) {
        List<GpuNode> result = new ArrayList<>();
        if (CollectionUtils.isNotEmpty(gpuNodes)) {
            List<WeightedGpuNode> weightedGpuNodes = new ArrayList<>();
            for (GpuNode gpuNode : gpuNodes) {
                WeightedGpuNode weightedGpuNode = new WeightedGpuNode();
                weightedGpuNode.setGpuNode(gpuNode);
                weightedGpuNode.setScore(gpuNode.getUsageCount().get());
                weightedGpuNodes.add(weightedGpuNode);
            }
            weightedGpuNodes.sort(Comparator.comparingInt(WeightedGpuNode::getScore));
            for (WeightedGpuNode weightedGpuNode : weightedGpuNodes) {
                result.add(weightedGpuNode.getGpuNode());
            }
        }
        return result;
    }

    private double computeScore(ContainerConfig config, Node node) {
        double cpuScore = (config.getCpu() + node.getNodeUsage().getUsedCpus()) /
                node.getResourceUsage().getTotalCpus();
        double memoryScore = (config.getMemory() + node.getNodeUsage().getUsedMemory()) /
                (1024.0 * node.getResourceUsage().getTotalMemory());
        return Math.round((cpuScore + memoryScore) * 100) / 100.0;
    }
}
